package cn.sciento.transfer.infra.rocketmq.utils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.Set;
import org.apache.rocketmq.acl.common.AclClientRPCHook;
import org.apache.rocketmq.acl.common.SessionCredentials;
import org.apache.rocketmq.common.TopicConfig;
import org.apache.rocketmq.common.admin.TopicStatsTable;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;
import org.apache.rocketmq.common.protocol.body.ClusterInfo;
import org.apache.rocketmq.common.protocol.body.SubscriptionGroupWrapper;
import org.apache.rocketmq.common.protocol.route.BrokerData;
import org.apache.rocketmq.common.protocol.route.TopicRouteData;
import org.apache.rocketmq.common.subscription.SubscriptionGroupConfig;
import org.apache.rocketmq.remoting.RPCHook;
import org.apache.rocketmq.tools.admin.DefaultMQAdminExt;
import org.apache.rocketmq.tools.admin.MQAdminExt;
import org.apache.rocketmq.tools.command.CommandUtil;
import org.apache.rocketmq.tools.command.SubCommandException;

public class RocketmqClientUtils {
  public static void createTopic(String nameSrc, String topic) throws SubCommandException {
    DefaultMQAdminExt adminExt = new DefaultMQAdminExt(getAclRPCHook());
    adminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    adminExt.setNamesrvAddr(nameSrc);
    try {
      adminExt.start();
      ClusterInfo clusterInfo = adminExt.examineBrokerClusterInfo();
      Set<String> brokerAddrList = new HashSet<>();
      for (String clusterName : clusterInfo.getClusterAddrTable().keySet())
        brokerAddrList.addAll(CommandUtil.fetchMasterAndSlaveAddrByClusterName((MQAdminExt)adminExt, clusterName)); 
      int randomIndex = (new Random()).nextInt(brokerAddrList.size());
      List<String> list = new ArrayList<>(brokerAddrList);
      String broker = list.get(randomIndex);
      TopicConfig topicConfig = new TopicConfig();
      topicConfig.setTopicName(topic);
      topicConfig.setWriteQueueNums(1);
      topicConfig.setReadQueueNums(1);
      topicConfig.setTopicSysFlag(0);
      adminExt.createAndUpdateTopicConfig(broker, topicConfig);
    } catch (Exception e) {
      throw new SubCommandException("command failed", e);
    } finally {
      adminExt.shutdown();
    } 
  }
  
  public static void deleteTopic(String nameSrc, String topic) throws SubCommandException {
    DefaultMQAdminExt adminExt = new DefaultMQAdminExt(getAclRPCHook());
    adminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    adminExt.setNamesrvAddr(nameSrc);
    try {
      adminExt.start();
      ClusterInfo clusterInfo = adminExt.examineBrokerClusterInfo();
      for (String clusterName : clusterInfo.getClusterAddrTable().keySet()) {
        Set<String> brokerAddressSet = CommandUtil.fetchMasterAndSlaveAddrByClusterName((MQAdminExt)adminExt, clusterName);
        adminExt.deleteTopicInBroker(brokerAddressSet, topic);
        Set<String> nameServerSet = null;
        if (adminExt.getNamesrvAddr() != null) {
          String[] ns = adminExt.getNamesrvAddr().trim().split(";");
          nameServerSet = new HashSet<>(Arrays.asList(ns));
        } 
        adminExt.deleteTopicInNameServer(nameServerSet, topic);
      } 
    } catch (Exception e) {
      throw new SubCommandException("command failed", e);
    } finally {
      adminExt.shutdown();
    } 
  }
  
  public static List<MessageQueue> queryQueueListByTopic(String nameSrc, String topic) throws SubCommandException {
    DefaultMQAdminExt defaultMQAdminExt = new DefaultMQAdminExt(getAclRPCHook());
    defaultMQAdminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    defaultMQAdminExt.setNamesrvAddr(nameSrc);
    try {
      defaultMQAdminExt.start();
      TopicStatsTable topicStatsTable = defaultMQAdminExt.examineTopicStats(topic);
      List<MessageQueue> mqList = new LinkedList<>();
      mqList.addAll(topicStatsTable.getOffsetTable().keySet());
      return mqList;
    } catch (Exception e) {
      throw new SubCommandException("command failed", e);
    } finally {
      defaultMQAdminExt.shutdown();
    } 
  }
  
  public static Long queryMessageOffsetByKey(String nameSrc, String topic, String eventId) throws SubCommandException {
    DefaultMQAdminExt defaultMQAdminExt = new DefaultMQAdminExt(getAclRPCHook());
    defaultMQAdminExt.setNamesrvAddr(nameSrc);
    defaultMQAdminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    try {
      defaultMQAdminExt.start();
      List<MessageExt> messageExts = defaultMQAdminExt.queryMessage(topic, eventId, 64, 0L, Long.MAX_VALUE).getMessageList();
      return ((MessageExt) messageExts.get(0)).getQueueOffset();
    } catch (Exception e) {
      throw new SubCommandException(" command failed", e);
    } finally {
      defaultMQAdminExt.shutdown();
    } 
  }
  
  public static void resetConsumerOffset(String nameSrc, String consumerGroup, String topic, MessageQueue messageQueue, Long offset) throws SubCommandException {
    DefaultMQAdminExt defaultMQAdminExt = new DefaultMQAdminExt(getAclRPCHook());
    defaultMQAdminExt.setNamesrvAddr(nameSrc);
    defaultMQAdminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    try {
      defaultMQAdminExt.start();
      TopicRouteData routeData = defaultMQAdminExt.examineTopicRouteInfo(topic);
      if (routeData.getBrokerDatas() != null) {
        List<String> brokerAddr = new ArrayList<>();
        for (BrokerData brokerData : routeData.getBrokerDatas())
          brokerAddr.addAll(brokerData.getBrokerAddrs().values()); 
        for (String broker : brokerAddr)
          defaultMQAdminExt.updateConsumeOffset(broker, consumerGroup, messageQueue, offset.longValue()); 
      } 
    } catch (Exception e) {
      throw new SubCommandException(" command failed", e);
    } finally {
      defaultMQAdminExt.shutdown();
    } 
  }
  
  public static List<String> getAllSubscriptionGroups(String nameSrc) throws SubCommandException {
    DefaultMQAdminExt defaultMQAdminExt = new DefaultMQAdminExt(getAclRPCHook());
    defaultMQAdminExt.setNamesrvAddr(nameSrc);
    defaultMQAdminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    try {
      defaultMQAdminExt.start();
      ClusterInfo clusterInfo = defaultMQAdminExt.examineBrokerClusterInfo();
      Set<String> brokerAddrList = new HashSet<>();
      for (String clusterName : clusterInfo.getClusterAddrTable().keySet())
        brokerAddrList.addAll(CommandUtil.fetchMasterAndSlaveAddrByClusterName((MQAdminExt)defaultMQAdminExt, clusterName)); 
      List<String> groups = new ArrayList<>();
      for (String brokerAddr : brokerAddrList) {
        SubscriptionGroupWrapper groupWrapper = defaultMQAdminExt.getAllSubscriptionGroup(brokerAddr, 20000L);
        groups.addAll(groupWrapper.getSubscriptionGroupTable().keySet());
      } 
      return groups;
    } catch (Exception e) {
      throw new SubCommandException(" command failed", e);
    } finally {
      defaultMQAdminExt.shutdown();
    } 
  }
  
  public static void createGroup(String nameSrv, String group) throws SubCommandException {
    DefaultMQAdminExt defaultMQAdminExt = new DefaultMQAdminExt(getAclRPCHook());
    defaultMQAdminExt.setNamesrvAddr(nameSrv);
    defaultMQAdminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    try {
      defaultMQAdminExt.start();
      ClusterInfo clusterInfo = defaultMQAdminExt.examineBrokerClusterInfo();
      Set<String> brokerAddrList = new HashSet<>();
      for (String clusterName : clusterInfo.getClusterAddrTable().keySet())
        brokerAddrList.addAll(CommandUtil.fetchMasterAndSlaveAddrByClusterName((MQAdminExt)defaultMQAdminExt, clusterName)); 
      Random random = new Random();
      int index = random.nextInt(brokerAddrList.size());
      SubscriptionGroupConfig subscriptionGroupConfig = new SubscriptionGroupConfig();
      subscriptionGroupConfig.setGroupName(group);
      defaultMQAdminExt.createAndUpdateSubscriptionGroupConfig(brokerAddrList.toArray()[index].toString(), subscriptionGroupConfig);
    } catch (Exception e) {
      throw new SubCommandException(" command failed", e);
    } finally {
      defaultMQAdminExt.shutdown();
    } 
  }
  
  public static void deleteGroup(String nameSrv, String group) throws SubCommandException {
    DefaultMQAdminExt defaultMQAdminExt = new DefaultMQAdminExt(getAclRPCHook());
    defaultMQAdminExt.setNamesrvAddr(nameSrv);
    defaultMQAdminExt.setInstanceName(Long.toString(System.currentTimeMillis()));
    try {
      defaultMQAdminExt.start();
      ClusterInfo clusterInfo = defaultMQAdminExt.examineBrokerClusterInfo();
      Set<String> brokerAddrList = new HashSet<>();
      for (String clusterName : clusterInfo.getClusterAddrTable().keySet())
        brokerAddrList.addAll(CommandUtil.fetchMasterAndSlaveAddrByClusterName((MQAdminExt)defaultMQAdminExt, clusterName)); 
      for (String addr : brokerAddrList)
        defaultMQAdminExt.deleteSubscriptionGroup(addr, group); 
    } catch (Exception e) {
      throw new SubCommandException(" command failed", e);
    } finally {
      defaultMQAdminExt.shutdown();
    } 
  }
  
  private static RPCHook getAclRPCHook() {
    return (RPCHook)new AclClientRPCHook(new SessionCredentials("DEFAULT", "DEFAULT"));
  }
}
